Skip to content

fix(dspark): reserve V2 padded Q scratch - #26

Merged
jasl merged 1 commit into
jasl:codex/ds4-sm120-min-enablefrom
alexbi29:codex/ds4-v2-q-scratch
Jul 3, 2026
Merged

fix(dspark): reserve V2 padded Q scratch#26
jasl merged 1 commit into
jasl:codex/ds4-sm120-min-enablefrom
alexbi29:codex/ds4-v2-q-scratch

Conversation

@alexbi29

@alexbi29 alexbi29 commented Jul 2, 2026

Copy link
Copy Markdown

Summary

  • Fixes the DSpark V2 runner OOM by removing the runtime padded-Q allocation from fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.
  • The fused op now writes into caller-owned q_out; DeepSeek V4 attention pre-reserves reusable padded-Q scratch during V2 profile/runtime.
  • Updates functionalization metadata and fused-kernel tests for the new mutating op signature.

Validation

Check Result
Live V2 startup Using V2 Model Runner, GPU KV cache 1,095,149 tokens, max concurrency 2.09x
Long-context concurrency No OOM / no HTTP 500; concurrency-8 at ~29.6k prompt tokens, 15/16 exact recall, one retrieval miss
TPS sanity avg 303.8 tok/s, median 302.0, range 297.4-311.5

@alexbi29

alexbi29 commented Jul 2, 2026

Copy link
Copy Markdown
Author

@jasl this should fix v2

@jasl
jasl merged commit 5b1b452 into jasl:codex/ds4-sm120-min-enable Jul 3, 2026
3 of 4 checks passed
@jasl

jasl commented Jul 3, 2026

Copy link
Copy Markdown
Owner

Thanks @alexbi29 — validated end-to-end on 2-node GB10 (TP=2, DSv4-Flash, FULL_AND_PIECEWISE) and merged. The fix is correct and does exactly what it says. Because the changed callsite (_fused_qnorm_rope_kv_insert) is shared with the V1 default runner + the MTP draft path, I ran the full V1/MTP matrix too, not just V2:

V1 (production default)

  • GSM8K 0.975 · issue19 instruction-following PASS · smoke coherent
  • Arthur needle-recall conc12 @~28k (batched FULL cudagraph): 24/24; @~434k conc1: 2/2
  • Memory: GPU KV cache 2,741,470 tok = pre-PR baseline exactly, no post-profile OOM (the persistent scratch does not shrink KV)

V1 + MTP2

  • GSM8K 0.953 · arthur conc12 24/24 · acceptance rate 0.91 / accept-len 2.57 (healthy)
  • decode 29.5 tok/s c1 / 48.8 c4, TPOT ~25–32 ms — perf-neutral (the change removes a per-forward alloc)

V2 runner — and this is the interesting part:

  • No OOM at conc8 / ~30k (the fix's purpose) ✓
  • Arthur recall conc8 @~28k: 16/16, GSM8K 0.96

That 16/16 is worth flagging. We had previously recorded V2 long-context recall collapsing at this exact config (we measured 3/16 @28k/conc8) and had written it off as an unfixable cudagraph-replay-class recall bug — the reason we'd stayed on V1. Your padded-Q OOM fix also resolves that collapse: the old "misses" were requests erroring under memory pressure, not a replay bug. So this reconciles your 15/16 vs our 3/16 — with the fix it's a clean 16/16 here, and V2 is back on the table for us. Nice catch.

Two small follow-ups I applied/noted on the merge (not blockers):

  1. Wired reserve_profile_scratch into the V1 runner too (it was V2-only). On V1 with enable_dbo=true, the ubatch-1 padded-Q scratch would otherwise be materialized after KV sizing (during ubatched capture) and charged against already-claimed KV → post-profile OOM risk. No-op for the DBO-off default (num_ubatches=1). Added as a follow-up commit so the shared branch is correct for all configs.
  2. tests/compile/passes/test_functionalization.py still registers the stale 3-arg torch.ops.vllm op, so CI only exercises the {1:q,2:k_cache} else-branch — the production _C {1:q_out,2:k_cache} branch (which the mutable-arg order in torch_bindings.cpp makes provably correct) is never tested. Worth updating the test op to the 4-tensor void _C schema to lock it against future schema drift.

Thanks again — merged into codex/ds4-sm120-min-enable (= this PR's base) + ds4-sm120-preview-dev.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants